import backtrader as bt
import pandas as pd
import numpy as np
from datetime import datetime


class LowVolStrategy(bt.Strategy):
    params = (
        ('bench_index', '000300.SH'),  # 基准指数
        ('n_stocks', 10),  # 持仓数量
        ('pe_threshold_high', 50),  # PE清仓阈值
        ('pe_threshold_low', 20),  # PE重启阈值
        ('crisis_days', 5),  # 极端行情判断周期
    )

    def __init__(self):
        self.stock_data = {}  # 存储各股票数据引用
        self.current_month = -1
        self.signal = True
        self.pe_history = []  # 存储历史PE数据（示例需替换真实数据）

        # 为每个数据创建波动率指标
        for d in self.datas[1:]:  # 假设datas[0]是指数数据
            self.stock_data[d._name] = {
                'stddev': bt.indicators.StandardDeviation(
                    d.close, period=20, plot=False
                )
            }

    def next(self):
        # 每月初重新选股
        if self.data.datetime.date(0).month != self.current_month:
            self.current_month = self.data.datetime.date(0).month
            self.rebalance_portfolio()

        # 执行风控逻辑
        self.execute_risk_control()

    def rebalance_portfolio(self):
        # 计算所有候选股票的波动率
        vol_ranking = []
        for d in self.datas[1:]:
            if len(d) < 20:  # 确保有足够数据
                continue
            vol = self.stock_data[d._name]['stddev'][0]
            vol_ranking.append((d, vol))

        # 选择波动率最低的n只股票
        vol_ranking.sort(key=lambda x: x[1])
        selected = [d for d, _ in vol_ranking[:self.params.n_stocks]]

        # 调整持仓
        self.adjust_positions(selected)

    def adjust_positions(self, selected):
        # 卖出不在候选列表中的持仓
        for d in self.getpositions():
            if d not in selected:
                self.close(d)

        # 等权重买入新标的
        weight = 1.0 / len(selected)
        for d in selected:
            self.order_target_percent(d, target=weight)

    def execute_risk_control(self):
        # 示例PE数据，需替换为真实数据源
        current_pe = self.get_current_pe()

        if current_pe > self.params.pe_threshold_high:
            self.signal = False
            self.close_all()
        elif current_pe < self.params.pe_threshold_low:
            self.signal = True

        if self.is_crisis():
            self.close_all()

    def get_current_pe(self):
        # 此处需要接入PE数据源，示例返回固定值
        return 30  # 替换为真实PE计算逻辑

    def is_crisis(self):
        # 判断最近5日指数表现
        index_data = self.datas[0]
        if len(index_data) < 5:
            return False

        returns = (index_data.close[0] / index_data.close[-5] - 1)
        volatility = np.std([index_data.close[i] / index_data.close[i - 1] - 1
                             for i in range(-5, 0)])
        return returns < -0.10 or volatility > 0.05

    def close_all(self):
        for d in self.getpositions():
            self.close(d)


class GenericCSV(bt.feeds.GenericCSVData):
    params = (
        ('dtformat', '%Y-%m-%d'),
        ('datetime', 0),
        ('open', 1),
        ('high', 2),
        ('low', 3),
        ('close', 4),
        ('volume', 5),
        ('openinterest', -1),
    )


def run_backtest():
    cerebro = bt.Cerebro()

    # 加载指数数据（假设第一个数据是基准指数）
    index_df = pd.read_csv('data/lgb_pricedata.csv')
    data = GenericCSV(dataname='path/to/index_data.csv')
    cerebro.adddata(data, name='index')

    # 加载个股数据
    stock_list = ['600000.SH', '000001.SZ']  # 示例股票列表
    for stock in stock_list:
        df = pd.read_csv(f'path/to/stock_data/{stock}.csv')
        data = GenericCSV(dataname=f'path/to/stock_data/{stock}.csv')
        cerebro.adddata(data, name=stock)

    cerebro.addstrategy(LowVolStrategy)
    cerebro.broker.setcash(1000000)
    cerebro.broker.setcommission(commission=0.0003)  # 佣金设置
    cerebro.addanalyzer(bt.analyzers.Returns, _name='returns')

    results = cerebro.run()
    print(f'最终资产价值: {cerebro.broker.getvalue():.2f}')


if __name__ == '__main__':
    run_backtest()